import requests
from bs4 import BeautifulSoup
from pymongo import MongoClient
from datetime import datetime
import os
import re

# 配置信息
mongoIP = 'mongodb://root:123456@127.0.0.1:27017'
CONFIG = {
    "target_url": "https://example.com/images",  # 替换为实际目标网站
    "headers": {
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
    },
    "mongodb": {
        "host": "127.0.0.1",
        "port": 27017,
        "db_name": "image_database",
        "username": "root",
        "password": "123456",
        "collection": "images"
    }
}

def connect_mongodb():
    """连接MongoDB数据库"""
    # client = MongoClient(
    #     host=CONFIG["mongodb"]["host"],
    #     port=CONFIG["mongodb"]["port"]
    # )
    client = MongoClient(
        mongoIP, connect=False
    )
    db = client[CONFIG["mongodb"]["db_name"]]
    return db[CONFIG["mongodb"]["collection"]]

def sanitize_filename(filename):
    """清理文件名中的非法字符"""
    return re.sub(r'[\\/*?:"<>|]', "", filename)

def save_to_mongodb(image_data, filename, source_url):
    """将图片保存到MongoDB"""
    collection = connect_mongodb()
    
    document = {
        "filename": filename,
        "image_data": image_data,
        "source_url": source_url,
        "created_at": datetime.now(),
        "metadata": {
            "file_size": len(image_data),
            "content_type": "image/jpeg"  # 根据实际情况修改
        }
    }
    
    try:
        result = collection.insert_one(document)
        print(f"Saved {filename} to MongoDB, ID: {result.inserted_id}")
        return True
    except Exception as e:
        print(f"Error saving to MongoDB: {str(e)}")
        return False

def download_image(img_url):
    """下载图片并返回二进制数据"""
    try:
        response = requests.get(
            img_url,
            headers=CONFIG["headers"],
            stream=True,
            timeout=10
        )
        if response.status_code == 200:
            return response.content
    except Exception as e:
        print(f"Error downloading {img_url}: {str(e)}")
    return None

def extract_images(url):
    """从网页中提取图片链接"""
    try:
        response = requests.get(url, headers=CONFIG["headers"], timeout=10)
        soup = BeautifulSoup(response.text, "html.parser")
        
        # 查找所有img标签
        img_tags = soup.find_all("img")
        
        # 提取图片URL（可能需要处理相对路径）
        image_urls = []
        for img in img_tags:
            src = img.get("src", "")
            if src.startswith("http"):
                image_urls.append(src)
            else:
                # 转换相对路径为绝对路径（根据网站结构可能需要调整）
                image_urls.append(f"{url}/{src}")
        
        return image_urls
    except Exception as e:
        print(f"Error extracting images: {str(e)}")
        return []

def main():
    # 获取目标网页中的图片链接
    image_urls = extract_images(CONFIG["target_url"])
    
    if not image_urls:
        print("No images found")
        return
    
    print(f"Found {len(image_urls)} images")
    
    for img_url in image_urls:
        # 下载图片
        image_data = download_image(img_url)
        if not image_data:
            continue
        
        # 生成文件名
        filename = sanitize_filename(os.path.basename(img_url))
        
        # 保存到MongoDB
        save_to_mongodb(image_data, filename, img_url)

if __name__ == "__main__":
    save_to_mongodb("", "1.png", "123")
    # main()